package grobit.wire

import grobit.PeerEngine
import grobit.torrent.ConnectionManager
import grobit.torrent.MetaInfo
import grobit.torrent.Sha1Data
import grobit.tracker.Peer
import grobit.tracker.PeerId
import grobit.tracker.TrackerRequest

import java.nio.ByteBuffer
import java.nio.channels.SocketChannel
import java.util.concurrent.atomic.AtomicInteger

/**
 * A class which can accept data and do something with it.
 */
interface DataHandler {
    void accept(ByteBuffer buffer)
}

class PeerConnection implements DataHandler {
    static final protocolName = 'BitTorrent protocol'
    static AtomicInteger idCounter = new AtomicInteger(0)

    ConnectionManager connectionManager
    int cId = idCounter.incrementAndGet()
    boolean amChoking = true
    boolean amInterested = false
    boolean peerChoking = true
    boolean peerInterested = false

    boolean busy = false

    Peer peer
    MetaInfo metaInfo
    SocketChannel socketChannel
    byte[] remotePeerId = new byte[20]
    DataHandler currentHandler
    BitSet peersPieces

    /**
     * This method sends the handshake and registers the connection with the PeerEngine to receive data - this
     * essentially makes the connection 'live' as the engine will execute callbacks in another thread.
     */
    void sendHandshake() {
        ByteBuffer byteBuffer = ByteBuffer.allocate(49 + protocolName.length())
        byteBuffer.put(protocolName.length() as byte)
        byteBuffer.put(protocolName.getBytes('ISO-8859-1'))
        byteBuffer.putLong(0)
        byteBuffer.put(metaInfo.infoHash.getBytes())
        byteBuffer.put(PeerId.get())
        byteBuffer.flip()
        socketChannel.write(byteBuffer)

        byte[] bitField = new byte[Math.ceil(metaInfo.pieces.size() / 8)]
        (0..<metaInfo.pieces.size()).each { pieceNum ->
            int pos = (pieceNum / 8) as int
            int bit = pieceNum % 8
            bitField[pos] = (bitField[pos] | (1<<bit))
        }
        ByteBuffer bitFieldBuf = ByteBuffer.allocate(bitField.length + 2)
        bitFieldBuf.put((bitField.length+1) as byte)
        bitFieldBuf.put(5 as byte)
        bitFieldBuf.put(bitField)
        bitFieldBuf.flip()
        dumpBuffer(bitFieldBuf)

        socketChannel.write(bitFieldBuf)
        println "$cId: Sent have all bitfields"

        currentHandler = new HandshakeHandler(parent: this)
        PeerEngine.instance.register(this)
    }

    void dumpBuffer(ByteBuffer buf) {
        print "$cId: Dump: "
        while(buf.remaining()) {
            print String.format("%02x ", (buf.get() & 0xFF))
        }
        println ''
        buf.flip()
    }

    /**
     * Requests a piece from the connection.
     *
     * @return The piece requested, or -1 if no piece was requested.
     */
    int requestNextPiece() {
        if (peerChoking || !amInterested) {
            println "$cId: Cannot request - wrong flags: $peerChoking $amInterested"
            return -1
        }

        if (!peersPieces || peersPieces.cardinality() == 0) {
            println "$cId: Cannot request - no pieces"
            return -1
        }

        busy = true
        int piece = peersPieces.nextSetBit(0)

        ByteBuffer buf = ByteBuffer.allocate(14)
        buf.put(13 as byte)
        buf.put(6 as byte)
        buf.putInt(piece)
        buf.putInt(0)
        buf.putInt(16384)
        buf.flip()
        socketChannel.write(buf)

        println "$cId: requested piece $piece"
        return piece
    }

    void processMessage(ByteBuffer buffer) {
        int messageType = (buffer.get() & 0xFF)
        println "$cId - Got full message - type: $messageType"
        switch(messageType) {
            case 0:
                println "$cId: Choke"
                peerChoking = true
                break
            case 1:
                println "$cId: Unchoke"
                peerChoking = false
                connectionManager.availabilityChanged(this)
                break
            case 2:
                println "$cId: Interested"
                peerInterested = true
                break
            case 3:
                println "$cId: Not-interested"
                peerInterested = false
                break
            case 4:
                int pieceIndex = buffer.getInt()
                int setSize = peersPieces.size()
                println "$cId: Peer has piece $pieceIndex of ${setSize}"
                peersPieces.set(pieceIndex)
                printBitfieldStats()
                if (peersPieces.cardinality() > 0) {
                    sendInterested()
                }
                connectionManager.availabilityChanged(this)
                break
            case 5:
                int bitPosition = 0;
                int setSize = peersPieces.size()
                while(buffer.remaining()) {
                    byte nextByte = buffer.get()
                    8.times {
                        if(nextByte && 0x80) {
                            peersPieces.set(bitPosition)
                        }
                        bitPosition++
                        nextByte << 1
                    }
                }
                println "$cId: Set bitfield ($bitPosition /$setSize)"
                printBitfieldStats()
                if (peersPieces.cardinality() > 0) {
                    sendInterested()
                }
                connectionManager.availabilityChanged(this)
                break
            case 6:
                println "$cId: Data requested: "
                break
            case 7:
                println "$cId: Piece data - " + buffer.remaining()
                break
            case 8:
                println "$cId: Cancel - unhandled"
                break
            case 9:
                println "$cId: Port Message - DHT not supported"
                break;
            default:
                println "$cId - unknown message type: $messageType"
                break
        }
    }

    void sendInterested() {
        if (amInterested) {
            return
        }

        ByteBuffer buffer = ByteBuffer.allocate(2)
        buffer.put(1 as byte)
        buffer.put(1 as byte)
        buffer.flip()
        socketChannel.write(buffer)
        println "$cId: Unchoke"
        amChoking = false

        buffer = ByteBuffer.allocate(2)
        buffer.put(1 as byte)
        buffer.put(2 as byte)
        buffer.flip()
        socketChannel.write(buffer)
        println "$cId: Registered interest"
        amInterested = true
    }

    void printBitfieldStats() {
        println "$cId: bitset: " + peersPieces.cardinality() + '/' + peersPieces.size()
    }

    void accept(ByteBuffer data) {
        currentHandler.accept(data)
    }

    void connect() {
        try {
            SocketChannel channel = SocketChannel.open(peer.address)
            println 'connection succeeded: ' + cId
            socketChannel = channel
        } catch(ConnectException x) {
            println 'Connection failed: ' + x.getMessage()
        }
        peersPieces = new BitSet(metaInfo.pieces.size())
    }

    boolean isConnected() {
        return socketChannel != null
    }

    void close() {
        if(socketChannel) {
            socketChannel.close()
        }
        socketChannel = null
    }

    @Override
    public String toString() {
        return "PeerConnection{" +
                "metaInfo=" + metaInfo +
                ", peer=" + peer +
                ", peerInterested=" + peerInterested +
                ", peerChoking=" + peerChoking +
                ", amInterested=" + amInterested +
                ", amChoking=" + amChoking +
                '}';
    }
}

class NullDataHandler implements DataHandler {
    PeerConnection parent

    @Override
    void accept(ByteBuffer buffer) {
        println "${parent.cId} NullHandler - accepted: " + buffer.remaining() + ' bytes'
        byte[] allData = new byte[buffer.remaining()]
        buffer.get(allData)
        println "${parent.cId} Got Data: " + TrackerRequest.asUrlString(new String(allData, 'ISO-8859-1'))
    }
}

class MessageHandler implements DataHandler {
    PeerConnection parent
    RequiredLengthBuffer lengthBuffer

    MessageHandler(int messageLength, PeerConnection parent) {
        lengthBuffer = new RequiredLengthBuffer(messageLength)
        this.parent = parent
    }

    @Override
    void accept(ByteBuffer buffer) {
        lengthBuffer.readFrom(buffer)
        if (!lengthBuffer.full) {
            return
        }

        parent.processMessage(lengthBuffer.internBuffer)

        parent.currentHandler = new MessageLengthHandler(parent: parent)
        parent.currentHandler.accept(buffer)
    }
}


class MessageLengthHandler implements DataHandler {
    PeerConnection parent
    RequiredLengthBuffer lengthBuffer = new RequiredLengthBuffer(4)

    @Override
    void accept(ByteBuffer buffer) {
        lengthBuffer.readFrom(buffer)
        if (! lengthBuffer.isFull()) {
            return
        }

        int length = lengthBuffer.internBuffer.getInt()
        if (length == 0) {
            println "$parent.cId - Keepalive received"
            parent.currentHandler = new MessageLengthHandler(parent: parent)
        } else {
            println "$parent.cId - Length $length received"
            parent.currentHandler = new MessageHandler(length, parent)
        }
        parent.currentHandler.accept(buffer)
    }
}

class HandshakeHandler implements DataHandler {

    PeerConnection parent
    RequiredLengthBuffer startByte = new RequiredLengthBuffer(1)
    RequiredLengthBuffer mainData = null
    int protocolLen

    @Override
    void accept(ByteBuffer buffer) {
        println "${parent.cId} HandshakeHandler - accepted: " + buffer.remaining() + ' bytes'
        if (mainData == null) {
            startByte.readFrom(buffer)
            if (!startByte.isFull()) {
                return
            }
            protocolLen = (0xFF & startByte.internBuffer.get())
            println 'protocol-len: ' + protocolLen
            mainData = new RequiredLengthBuffer(48 + protocolLen)
        }

        mainData.readFrom(buffer)
        if( ! mainData.full) {
            return
        }

        byte[] protocolName = new byte[protocolLen]
        mainData.internBuffer.get(protocolName)

        long optionBits = mainData.internBuffer.getLong()
        println "$parent.cId - Option bits $optionBits"

        ByteBuffer byteBuffer = mainData.internBuffer
        byte[] remoteHash = new byte[20]
        byteBuffer.get(remoteHash)
        Sha1Data sha1Data = new Sha1Data(remoteHash)
        if (sha1Data != parent.metaInfo.infoHash) {
            println "${parent.cId} SHA-1 incorrect: $sha1Data ($parent.metaInfo.infoHash)"
        } else {
            println "${parent.cId} SHA-1 OK - $sha1Data"
        }

        parent.remotePeerId = new byte[20]
        byteBuffer.get(parent.remotePeerId)
        println "${parent.cId} Activating: " + (TrackerRequest.asUrlString(new String(parent.remotePeerId, 'ISO-8859-1')))

        parent.currentHandler = new MessageLengthHandler(parent: parent)
        parent.currentHandler.accept(buffer)

    }
}

class RequiredLengthBuffer {
    ByteBuffer internBuffer
    boolean full = false

    RequiredLengthBuffer(int length) {
        internBuffer = ByteBuffer.allocate(length)
    }

    void readFrom(ByteBuffer incoming) {
        if (full) {
            return
        }
        int remaining = incoming.limit() - incoming.position()
        int willTake = Math.min(remaining, remainingRequired())
        byte[] transfer = new byte[willTake]
        incoming.get(transfer)
        internBuffer.put(transfer)
        if(remainingRequired() == 0) {
            full = true
            internBuffer.flip()
        }
    }

    int remainingRequired() {
        return full ? 0 : internBuffer.capacity() - internBuffer.position()
    }

    boolean isFull() {
        return full
    }
}